# -*- coding: utf-8 -*-
import os
import sys
import argparse
import shutil


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--dataset", help="Folder with data", required=True)
    return parser.parse_args()
    

"""
Use this script for separation one big set of files to folders (for tree_leaves_test, for example
"""
if __name__ == '__main__':
    args = parse_args()
        
    for root, dirs, files in os.walk(args.dataset):
        filenames = [os.path.join(root, name) for name in files]
        folder_name = os.path.split(root)[-1].strip(' .0123456789')  #Clear digits and leading\trailing spaces
        filenames = [x for x in filenames if x.split('.')[-1].lower() in ['png', 'tiff', 'jpg']]
        if not filenames:
            continue
            
        for file in files:
            class_name = file[:-5].strip(' .0123456789_')
            print class_name
            src = os.path.join(root, file)
            try:
                os.makedirs(os.path.join(root, class_name))
            except:
                pass
            shutil.move(src, os.path.join(root, class_name, file))